#include <jni.h>
#include <string>
#include "android/log.h"
#include "turbojpeg.h"

#define CLASS_JAVA "top/broncho/libjpeg/JpegUtils"

#define LOG_TAG "JpegUtils"
#define LOGD(...) __android_log_print(ANDROID_LOG_DEBUG, LOG_TAG, __VA_ARGS__)
#define LOGE(...) __android_log_print(ANDROID_LOG_ERROR, LOG_TAG, __VA_ARGS__)

extern "C" {
#include "jpeglib.h"
}

#ifdef __cplusplus
extern "C" {
#endif


int yuv2jpeg(unsigned char *yuv_buffer, int width, int height, int subsamp,
             unsigned char **jpeg_buffer, unsigned long *jpeg_size, int quality) {
    int flags = 0;
    int padding = 1; //must be 1 or 4

    tjhandle handle = tjInitCompress();

    unsigned long yuv_size = width * height * 3 / 2;
    unsigned long need_size = tjBufSizeYUV2(width, padding, height, subsamp);
    if (need_size != yuv_size) {
        LOGE("we detect yuv size: %ld, but you give: %ld, check again.", need_size, yuv_size);
        return 0;
    }

    int ret = tjCompressFromYUV(handle, yuv_buffer, width, padding, height, subsamp, jpeg_buffer,
                                jpeg_size, quality, flags);
    if (ret < 0) {
        LOGE("compress to jpeg failed: %s.", tjGetErrorStr());
    }

    tjDestroy(handle);

    return ret;
}

//NV21->YUV420P
void yuv420sp_to_yuv420p(unsigned char *yuv420sp, unsigned char *yuv420p, int width, int height) {
    int i, j;
    int y_size = width * height;

    unsigned char *y = yuv420sp;
    unsigned char *uv = yuv420sp + y_size;

    unsigned char *y_tmp = yuv420p;
    unsigned char *u_tmp = yuv420p + y_size;
    unsigned char *v_tmp = yuv420p + y_size * 5 / 4;

    // y
    memcpy(y_tmp, y, y_size);

    // u
    for (j = 0, i = 0; j < y_size / 2; j += 2, i++) {
        v_tmp[i] = uv[j];
        u_tmp[i] = uv[j + 1];
    }
}

jbyteArray Jni_NV21ToJPEG(JNIEnv *env, jclass clazz, jbyteArray nv21, jint width,
                          jint height) {
    LOGD("Jni_NV21ToJPEG: %dx%d", width, height);
    jbyte *yuv_buffer = env->GetByteArrayElements(nv21, JNI_FALSE);
    auto yuv_size = width * height * 3 / 2;
    auto *i420 = malloc(yuv_size);

    yuv420sp_to_yuv420p((unsigned char *) yuv_buffer, (unsigned char *) i420, width, height);

    unsigned char *jpeg_buffer;
    unsigned long jpeg_size = 0;
    yuv2jpeg((unsigned char *) i420, width, height, TJSAMP_420, &jpeg_buffer, &jpeg_size,
             100); //50:quality
    LOGD("Jni_NV21ToJPEG: jpeg_size=%ld", jpeg_size);

    jbyteArray array = env->NewByteArray(jpeg_size);
    LOGD("Jni_NV21ToJPEG: Sizeof buf= %ld", sizeof(jpeg_buffer));
    env->SetByteArrayRegion(array, 0, jpeg_size, (jbyte *) jpeg_buffer);

    free(i420);
    free(jpeg_buffer);
    return array;
}

static JNINativeMethod g_methods[] = {
        {"NV21ToJPEG", "([BII)[B", (void *) Jni_NV21ToJPEG}
};

JNIEXPORT jint JNI_OnLoad(JavaVM *vm, void *reserved) {
    JNIEnv *env = nullptr;

    if (vm->GetEnv((void **) &env, JNI_VERSION_1_4) != JNI_OK) {
        return JNI_ERR;
    }
    assert(env != nullptr);
    jclass clazz = env->FindClass(CLASS_JAVA);
    env->RegisterNatives(clazz, g_methods, (int) (sizeof(g_methods) / sizeof((g_methods)[0])));

    return JNI_VERSION_1_4;
}

JNIEXPORT void JNI_OnUnload(JavaVM *jvm, void *reserved) {
    JNIEnv *env = nullptr;

    if (jvm->GetEnv((void **) &env, JNI_VERSION_1_4) != JNI_OK) {
        return;
    }
    jclass clazz = env->FindClass(CLASS_JAVA);
    env->UnregisterNatives(clazz);
}


#ifdef __cplusplus
}
#endif